# SIPG for 2D Poisson with pure Dirichlet BC in dolfinx
# u_exact = exp(-x - y^2); f = (1 - 4*y^2) * u_exact
from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
import ufl
from ufl import dx, ds, dS, grad, inner, avg, jump
from dolfinx import mesh, fem, io, la
import basix
from dolfinx.fem.petsc import assemble_matrix, assemble_vector

# -----------------------
# 用户可调参数
# -----------------------
nx, ny = 32, 32              # 网格划分
degree = 1                   # DG 次数 p
C_pen = 20.0                 # 罚参数基数
cell_type = mesh.CellType.triangle  # 或 mesh.CellType.quadrilateral

# -----------------------
# 网格与函数空间
# -----------------------
domain = mesh.create_unit_square(MPI.COMM_WORLD, nx, ny, cell_type=cell_type)
V = fem.functionspace(domain, ("DG", degree))

u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)

x = ufl.SpatialCoordinate(domain)
# 解析解与右端
u_exact = ufl.exp(-x[0] - x[1]*x[1])
f_expr = (1.0 - 4.0*x[1]*x[1]) * u_exact

# 边界法向、单元尺度
n = ufl.FacetNormal(domain)
h = ufl.CellDiameter(domain)
h_avg = (h('+') + h('-')) / 2.0
sigma_int = C_pen * (degree + 1)**2 / h_avg     # 内部面惩罚
sigma_bnd = C_pen * (degree + 1)**2 / h         # 边界面惩罚

# -----------------------
# SIPG 变分形式
# -----------------------
# 体项
a_vol = inner(grad(u), grad(v)) * dx

# 内部面项（对称内罚）
a_int = (
    - inner(avg(grad(u)), jump(v, n)) * dS
    - inner(avg(grad(v)), jump(u, n)) * dS
    + inner(jump(u, n), jump(v, n)) * sigma_int * dS
)

# 纯 Dirichlet 边界的 Nitsche/SIPG 项
a_bnd = (
    - inner(grad(u), n) * v * ds
    - inner(grad(v), n) * u * ds
    + sigma_bnd * u * v * ds
)

a = a_vol + a_int + a_bnd

# 右端（含边界贡献）
L = (
    f_expr * v * dx
    - inner(grad(v), n) * u_exact * ds
    + sigma_bnd * u_exact * v * ds
)

# -----------------------
# 组装与求解
# -----------------------
A = assemble_matrix(fem.form(a))
A.assemble()
b = assemble_vector(fem.form(L))

# 求解 Ax=b
ksp = PETSc.KSP().create(MPI.COMM_WORLD)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")

x_sol = A.createVecRight()
uh = fem.Function(V)

ksp.solve(b, x_sol)
uh.x.array[:] = x_sol.array

# -----------------------
# 误差评估
# -----------------------
# L2 误差
e = uh - u_exact
L2_err = np.sqrt(fem.assemble_scalar(fem.form(e*e*dx)))
# H1 半范误差（基于梯度）
H1s_err = np.sqrt(fem.assemble_scalar(fem.form(inner(grad(e), grad(e))*dx)))

# DG-“能量”型范数（常用于IP方法）
jump_term = fem.assemble_scalar(fem.form(inner(jump(e, n), jump(e, n)) * sigma_int * dS))
boundary_term = fem.assemble_scalar(fem.form(sigma_bnd * e*e * ds))
DG_energy = np.sqrt(
    fem.assemble_scalar(fem.form(inner(grad(e), grad(e))*dx))
    + jump_term
    + boundary_term
)

if MPI.COMM_WORLD.rank == 0:
    print(f"L2 error       = {L2_err:.6e}")
    print(f"H1 seminorm    = {H1s_err:.6e}")
    print(f"DG energy norm = {DG_energy:.6e}")

